import pymysql
from config import config

# 连接数据库并添加avatar字段到user表
def add_avatar_column():
    connection = None
    try:
        connection = pymysql.connect(
            host=config.MYSQL_HOST,
            port=config.MYSQL_PORT,
            user=config.MYSQL_USER,
            password=config.MYSQL_PASSWORD,
            database=config.MYSQL_DATABASE,
            charset='utf8mb4',
            cursorclass=pymysql.cursors.DictCursor
        )
        
        with connection.cursor() as cursor:
            # 检查avatar字段是否已存在
            cursor.execute("SHOW COLUMNS FROM user LIKE 'avatar'")
            result = cursor.fetchone()
            
            if not result:
                # 添加avatar字段
                cursor.execute("ALTER TABLE user ADD COLUMN avatar VARCHAR(500) NULL DEFAULT NULL COMMENT '用户头像URL'")
                connection.commit()
                print("Avatar column added successfully!")
            else:
                print("Avatar column already exists!")
                
    except Exception as e:
        print(f"Error: {e}")
    finally:
        if connection:
            connection.close()

if __name__ == "__main__":
    add_avatar_column()